from langchain.tools import BaseTool
from pydantic import BaseModel, Field, ConfigDict
from typing import Any, Type, List, Dict
import os
import requests
from dotenv import load_dotenv
from requests.exceptions import Timeout, ConnectionError, HTTPError

load_dotenv()
SERP_API_KEY = os.getenv("SERPAPI_API_KEY")
SERP_API_URL = os.getenv("SERP_API_URL")  # SerpAPI基础接口
SEARCH_TIMEOUT = 10  # 超时时间（秒），避免长期阻塞


class SearchToolInput(BaseModel):
    """搜索工具输入参数模型（Pydantic校验）"""
    query: str = Field(
        ...,  # 必传参数
        description="搜索关键词，需精准、含关键信息（如“贵州茅台2024年Q4财报 营收 净利润”“白酒行业2025政策”）"
    )
    result_count: int = Field(
        3,  # 默认返回3条结果（避免信息冗余）
        ge=1,  # 最少1条
        le=5,  # 最多5条（防止API调用成本过高）
        description="返回的搜索结果数量，建议1-5条"
    )


class SearchTool(BaseTool):
    """联网搜索工具（基于SerpAPI），用于获取本地库（Neo4j/文档）中没有的实时/最新财务信息"""
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **kwargs: Any):
        super().__init__(
            name="get_search_tool",  # Agent调用工具的名称（必须与Prompt中一致）
            description="""
            用于联网获取实时/最新信息，仅在以下场景调用：
            1. 用户需求是真实数据，且Neo4jTool、DocumentTool均返回“未找到数据”；
            2. 用户需求含“最新/实时/2025”等时效性关键词（如“2025白酒行业政策”“贵州茅台最新股价”）；
            3. 需求涉及行业动态、市场数据等本地库无法覆盖的信息（如“白酒行业2024年毛利率排名”）。

            调用约束：
            - 关键词需精准（含公司名、时间、数据类型，如“贵州茅台2024年Q4财报 净利润”）；
            - 若返回“无有效结果”，需停止搜索，返回不知道；
            - 不可用于查询与财务无关的信息（如“天气”“新闻八卦”）。
            """,
            args_schema=SearchToolInput,  # 绑定输入参数模型（自动校验）
            **kwargs
        )

    def _run(self, query: str, result_count: int = 3) -> str:
        """
        核心执行方法：调用SerpAPI搜索，返回结构化结果
        :param query: 搜索关键词
        :param result_count: 返回结果数量
        :return: 结构化搜索结果（文本格式，含标题、摘要、链接）
        """
        # 1. 校验API密钥（无密钥时返回错误提示）
        if not SERP_API_KEY:
            return "⚠️ 联网搜索失败：未配置有效SerpAPI密钥（需从https://serpapi.com/获取并设置环境变量SERP_API_KEY）"

        # 2. 构造搜索参数（SerpAPI要求的格式，支持Google搜索）
        params = {
            "q": query,  # 搜索关键词
            "api_key": SERP_API_KEY,  # API密钥
            "engine": "baidu",  # 搜索引擎（可选bing，需改参数为"engine":"bing"）
            "hl": "zh-CN",  # 语言（中文）
            "gl": "cn",  # 地区（中国）
            "num": result_count  # 返回结果数量
        }

        try:
            # 3. 发送搜索请求（带超时控制）
            response = requests.get(
                url=SERP_API_URL,
                params=params,
                timeout=SEARCH_TIMEOUT
            )
            response.raise_for_status()  # 触发HTTP错误（如401密钥无效、429请求超限）
            search_result = response.json()

            # 4. 提取有效结果（过滤广告、非文本内容）
            organic_results = search_result.get("organic_results", [])  # 自然搜索结果
            if not organic_results:
                return f"❌ 未获取到与「{query}」相关的有效搜索结果（可能数据未公开或关键词不准确）"

            # 5. 结构化结果（整理为AI可理解的文本）
            structured_result = f"✅ 联网搜索到「{query}」相关结果（共{min(len(organic_results), result_count)}条）：\n"
            for idx, result in enumerate(organic_results[:result_count], 1):
                # 提取结果中的核心字段（标题、摘要、链接，无字段时用“-”填充）
                title = result.get("title", "无标题")
                snippet = result.get("snippet", "无摘要")
                link = result.get("link", "无链接")
                # 拼接单条结果
                structured_result += f"""
                {idx}. 标题：{title}
                   摘要：{snippet}
                   链接：{link}
                """

            # 6. 补充提示（帮助AI判断是否使用结果）
            structured_result += "\n💡 提示：若结果包含明确财务数据（如营收、净利润），可基于此生成图表；若无明确数据，需生成虚拟数据。"
            return structured_result

        # 7. 捕获异常（分场景返回提示）
        except Timeout:
            return f"⚠️ 联网搜索超时（超过{SEARCH_TIMEOUT}秒），可能是网络问题或API响应缓慢，建议重试或改用虚拟数据"
        except ConnectionError:
            return "⚠️ 联网搜索失败：网络连接异常（请检查网络状态）"
        except HTTPError as e:
            if e.response.status_code == 401:
                return "⚠️ 联网搜索失败：API密钥无效（请检查SERP_API_KEY是否正确）"
            elif e.response.status_code == 429:
                return "⚠️ 联网搜索失败：API请求次数超限（SerpAPI免费额度不足，需升级套餐或次日重试）"
            else:
                return f"⚠️ 联网搜索HTTP错误：{str(e)}"
        except Exception as e:
            return f"⚠️ 联网搜索未知错误：{str(e)}"

    def _arun(self, query: str, result_count: int = 3) -> None:
        """异步执行方法（LangChain Tool要求实现，若不使用异步可直接抛出异常）"""
        raise NotImplementedError("SearchTool 不支持异步调用，请使用同步方法 _run")


if __name__ == "__main__":
    # 初始化工具
    search_tool = SearchTool()
    # 测试搜索（需先配置SERP_API_KEY环境变量）
    result = search_tool._run(
        query="贵州茅台2024年Q4财报 营收 净利润",
        result_count=2
    )
    # 打印结果
    print("搜索结果：")
    print(result)
